import os

import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision.transforms as transforms
from skimage.metrics import structural_similarity as ssim
from sklearn.metrics import f1_score, jaccard_score
from torch.utils.data import DataLoader

from Dataloader import CustomerDataLoader
from model import MedT


# 定义评价指标
def evaluate_metrics(test_Y, pred_Y):
    mse = torch.mean((test_Y - pred_Y) ** 2).item()  # MSE 计算
    test_Y_np = test_Y.squeeze(1).cpu().numpy()  # (B, H, W) 变为 numpy 数组
    pred_Y_np = pred_Y.squeeze(1).cpu().numpy()  # (B, H, W) 变为 numpy 数组

    # SSIM 计算：使用 np.min/max 来确保数据范围正确
    ssim_score = np.mean([
        ssim(test_Y_np[i], pred_Y_np[i], data_range=1.0)  # 假设数据范围为 [0, 1]
        for i in range(test_Y_np.shape[0])
    ])

    # 二值化处理
    binary_test_Y = (test_Y_np > 0.5).astype(int)
    binary_pred_Y = (pred_Y_np > 0.5).astype(int)

    # Dice 系数
    dice_coefficient = np.mean([
        f1_score(binary_test_Y[i].flatten(), binary_pred_Y[i].flatten())
        for i in range(binary_test_Y.shape[0])
    ])

    # 交并比 IoU
    iou = np.mean([
        jaccard_score(binary_test_Y[i].flatten(), binary_pred_Y[i].flatten())
        for i in range(binary_test_Y.shape[0])
    ])

    return mse, ssim_score, dice_coefficient, iou


# 测试函数
def test(model, dataloader, device):
    model.eval()  # 设置模型为评估模式
    all_metrics = {
        'mse': [],
        'ssim': [],
        'dice': [],
        'iou': []
    }
    total_images = 0  # 用于累加处理的图片数量
    with torch.no_grad():
        for i, [test_X, test_Y] in enumerate(dataloader):
            batch_size = test_X.size(0)
            total_images += batch_size

            test_X = test_X.float().to(device)
            test_Y = test_Y.float().to(device)

            # 前向传播进行预测
            pred_Y = model(test_X)

            # 如果模型输出是列表，选择最后一个输出
            if isinstance(pred_Y, list):
                pred_Y = pred_Y[-1]

            # 计算评价指标
            mse, ssim_score, dice_coefficient, iou = evaluate_metrics(test_Y, pred_Y)
            all_metrics['mse'].append(mse)
            all_metrics['ssim'].append(ssim_score)
            all_metrics['dice'].append(dice_coefficient)
            all_metrics['iou'].append(iou)

            print(f"Sample {i + 1}: MSE={mse:.4f}, SSIM={ssim_score:.4f}, Dice={dice_coefficient:.4f}, IoU={iou:.4f}")

            # 可视化部分
            if i <= 30:  # 只展示前 30 个测试样本

                # 转换为 NumPy 格式
                test_X_np = test_X[0].permute(1, 2, 0).cpu().numpy()
                test_Y_np = test_Y[0].squeeze(0).cpu().numpy()
                pred_Y_np = pred_Y[0].squeeze(0).cpu().numpy()
                pred_Y_np = (pred_Y_np > 0.5).astype(int)
                # 创建画布，并添加子图
                plt.figure(figsize=(12, 4))

                # 显示输入图像
                plt.subplot(1, 3, 1)
                plt.imshow(test_X_np)
                plt.title("Input Image")
                plt.xticks([])  # 去掉横坐标值
                plt.yticks([])  # 去掉纵坐标值
                plt.axis('off')

                # 显示真实标签
                plt.subplot(1, 3, 2)
                plt.imshow(test_Y_np, cmap='gray')
                plt.title("True Mask")
                plt.xticks([])  # 去掉横坐标值
                plt.yticks([])  # 去掉纵坐标值
                plt.axis('off')

                # 显示模型的预测结果
                plt.subplot(1, 3, 3)
                plt.imshow(pred_Y_np, cmap='gray')
                plt.title("Predicted Mask")
                plt.axis('off')
                plt.xticks([])  # 去掉横坐标值
                plt.yticks([])  # 去掉纵坐标值
                plt.show()

    # 计算平均指标
    avg_metrics = {metric: np.mean(values) for metric, values in all_metrics.items()}
    print("\nAverage Metrics:")
    print(f"Mean Squared Error: {avg_metrics['mse']:.4f}")
    print(f"Structural Similarity Index: {avg_metrics['ssim']:.4f}")
    print(f"Dice Coefficient: {avg_metrics['dice']:.4f}")
    print(f"Intersection over Union: {avg_metrics['iou']:.4f}")
    # 打印总处理的图片数量
    print(f"Total images tested: {total_images}")


if __name__ == '__main__':

    # 创建测试数据集
    img_dir = "/Volumes/For_Mac/Download/GlasDateset/test/images"  # 图片文件夹路径
    mask_dir = "/Volumes/For_Mac/Download/GlasDateset/test/masks"  # mask文件夹路径

    # 自定义的transform，包含resize和转换为tensor操作
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # 调整图像大小为224x224
        transforms.ToTensor(),  # 转换为Tensor格式
    ])
    dataset = CustomerDataLoader(img_dir, mask_dir, transform=transform)

    # 创建数据加载器
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2)

    # 初始化模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = MedT().to(device)

    # 加载训练好的模型参数
    model_path = '../pt_file/MedT_epoch_400.pt'  # 确保这是你保存模型的正确路径
    if os.path.exists(model_path):
        model.load_state_dict(torch.load(model_path, map_location=device))
        print(f"Model loaded from {model_path}")
    else:
        print(f"Model file {model_path} does not exist.")
        exit()

    # 开始测试
    test(model, dataloader, device)
